import os
import sys
import hashlib
import sqlite3
import argparse
from datetime import datetime
from tqdm import tqdm

def calculate_file_hash(file_path):
    """计算文件的SHA256哈希值"""
    hasher = hashlib.sha256()
    with open(file_path, 'rb') as f:
        buf = f.read()
        hasher.update(buf)
    return hasher.hexdigest()

def create_db_table(conn):
    """创建表结构"""
    cursor = conn.cursor()
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS video_hash (
            id INTEGER PRIMARY KEY,
            file_path TEXT NOT NULL,
            file_name TEXT NOT NULL,
            file_hash TEXT NOT NULL,
            create_time TEXT NOT NULL
        )
    ''')
    conn.commit()

def insert_into_db(conn, file_path, file_name, file_hash):
    """插入数据到数据库"""
    cursor = conn.cursor()
    
    cursor.execute('SELECT 1 FROM video_hash WHERE file_hash=?', (file_hash,))
    now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    if cursor.fetchone():
        return now
        
    cursor.execute('''
        INSERT INTO video_hash (file_path, file_name, file_hash, create_time) VALUES (?, ?, ?, ?)
    ''', (file_path, file_name, file_hash, now))
    conn.commit()
    return now  # 返回插入时间

def scan_directory_and_store_hashes(root_dir, scan_dir, db_conn):

    # 计算总文件数，用于进度条
    total_files = sum(len(files) for _, _, files in os.walk(scan_dir))
    with tqdm(total=total_files, desc='Processing files', unit='file') as pbar:
        for root, dirs, files in os.walk(scan_dir):
            for name in files:
                file_path = os.path.join(root, name)
                relative_path = os.path.relpath(file_path, root_dir)
                file_hash = calculate_file_hash(file_path)
                insert_time = insert_into_db(db_conn, relative_path, name, file_hash)
                pbar.update(1)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Scan files under the specified directory and compute their hashes to be stored in an SQLite database")
    parser.add_argument('--db', required=True, help='Directory where the SQLite database is located')
    parser.add_argument('--scan-dir', required=True, help='Directory to scan')

    args = parser.parse_args()
    db_path_absolute = os.path.abspath(args.db)
    scan_dir_absolute = os.path.abspath(args.scan_dir)
    # 检查 --db 和 --scan-dir 目录是否存在
    if not os.path.isdir(db_path_absolute):
        print(f"Error: 目录 {db_path_absolute} 不存在")
        exit(1)

    if not os.path.isdir(scan_dir_absolute):
        print(f"Error: 目录 {scan_dir_absolute} 不存在")
        exit(1)
    # 检查 --scan-dir 是否以 --db 开头，如果不是，则打印错误信息并退出
    if not scan_dir_absolute.startswith(db_path_absolute):
        print("Error: --scan-dir 参数的值必须以 --db 参数的值开头")
        exit(1)

    db_path = os.path.join(args.db, 'file_hash.db')
    video_dir = args.scan_dir

    # 创建/连接数据库
    conn = sqlite3.connect(db_path)
    
    # 创建表
    create_db_table(conn)

    # 扫描目录并存储哈希值
    scan_directory_and_store_hashes(args.db, video_dir, conn)

    # 关闭数据库连接
    conn.close()